//! Function overload resolver: annotates call expressions with resolved targets.

use crate::ast::*;
use crate::error::{CompilerError, Result};
use crate::utils::mangling;
use std::collections::HashMap;

use super::symbol_table::SymbolTable;

/// Resolves overloaded function calls using declared parameter types.
pub struct OverloadResolver<'a> {
    symbol_table: &'a SymbolTable,
    variable_env: Vec<HashMap<String, Type>>,
}

impl<'a> OverloadResolver<'a> {
    /// Create a new resolver.
    pub fn new(symbol_table: &'a SymbolTable) -> Self {
        Self {
            symbol_table,
            variable_env: Vec::new(),
        }
    }

    /// Resolve all overloads in the program.
    pub fn resolve_program(&mut self, program: &mut Program) -> Result<()> {
        for declaration in &mut program.declarations {
            self.resolve_declaration(declaration)?;
        }
        Ok(())
    }

    fn resolve_declaration(&mut self, declaration: &mut Declaration) -> Result<()> {
        match declaration {
            Declaration::Function(func) => self.resolve_function(func),
            Declaration::Struct(struct_decl) => {
                for method in &mut struct_decl.methods {
                    self.resolve_function(method)?;
                }
                Ok(())
            }
            Declaration::Impl(impl_decl) => {
                for method in &mut impl_decl.methods {
                    self.resolve_function(method)?;
                }
                Ok(())
            }
            Declaration::TraitImpl(trait_impl) => {
                for method in &mut trait_impl.methods {
                    self.resolve_function(method)?;
                }
                Ok(())
            }
            _ => Ok(()),
        }
    }

    fn resolve_function(&mut self, func: &mut FunctionDecl) -> Result<()> {
        if let Some(body) = &mut func.body {
            self.enter_scope();
            for param in &func.parameters {
                self.bind_variable(&param.name, param.param_type.clone());
            }
            for statement in &mut body.statements {
                self.resolve_statement(statement)?;
            }
            self.exit_scope();
        }
        Ok(())
    }

    fn resolve_statement(&mut self, statement: &mut Statement) -> Result<()> {
        match statement {
            Statement::Expression(expr) => {
                self.resolve_expression(expr)?;
            }
            Statement::VariableDecl(var_decl) => {
                if let Some(init) = &mut var_decl.initializer {
                    self.resolve_expression(init)?;
                }
                if let Some(var_type) = &var_decl.var_type {
                    self.bind_variable(&var_decl.name, var_type.clone());
                }
            }
            Statement::ConstantDecl(const_decl) => {
                self.resolve_expression(&mut const_decl.value)?;
                if let Some(const_type) = &const_decl.const_type {
                    self.bind_variable(&const_decl.name, const_type.clone());
                }
            }
            Statement::Assignment(assign) => {
                self.resolve_expression(&mut assign.value)?;
            }
            Statement::Return(ret) => {
                if let Some(value) = &mut ret.value {
                    self.resolve_expression(value)?;
                }
            }
            Statement::Block(block_stmt) => {
                self.enter_scope();
                for stmt in &mut block_stmt.statements {
                    self.resolve_statement(stmt)?;
                }
                self.exit_scope();
            }
            Statement::If(if_stmt) => {
                self.resolve_expression(&mut if_stmt.condition)?;
                self.resolve_statement(&mut if_stmt.then_branch)?;
                if let Some(else_branch) = &mut if_stmt.else_branch {
                    self.resolve_statement(else_branch)?;
                }
            }
            Statement::Match(match_stmt) => {
                self.resolve_expression(&mut match_stmt.expr)?;
                for arm in &mut match_stmt.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.resolve_expression(guard)?;
                    }
                    self.resolve_statement(&mut arm.body)?;
                }
            }
            Statement::While(while_stmt) => {
                self.resolve_expression(&mut while_stmt.condition)?;
                self.resolve_statement(&mut while_stmt.body)?;
            }
            Statement::For(for_stmt) => {
                match &mut for_stmt.variant {
                    ForVariant::Traditional {
                        initializer,
                        condition,
                        increment,
                    } => {
                        if let Some(init) = initializer {
                            self.resolve_statement(init)?;
                        }
                        if let Some(cond) = condition {
                            self.resolve_expression(cond)?;
                        }
                        if let Some(inc) = increment {
                            self.resolve_statement(inc)?;
                        }
                    }
                    ForVariant::Condition { condition } => {
                        self.resolve_expression(condition)?;
                    }
                    ForVariant::Range {
                        variable,
                        variable_type,
                        iterable,
                    } => {
                        self.resolve_expression(iterable)?;
                        self.enter_scope();
                        self.bind_variable(variable, variable_type.clone());
                        self.resolve_statement(&mut for_stmt.body)?;
                        self.exit_scope();
                        return Ok(());
                    }
                }
                self.resolve_statement(&mut for_stmt.body)?;
            }
            Statement::Defer(defer_stmt) => {
                self.resolve_expression(&mut defer_stmt.expr)?;
            }
            Statement::Go(go_stmt) => {
                self.resolve_call(&mut go_stmt.call)?;
            }
            Statement::Select(select_stmt) => {
                for case in &mut select_stmt.cases {
                    self.resolve_expression(&mut case.channel)?;
                    self.resolve_statement(&mut case.body)?;
                }
                if let Some(default_case) = &mut select_stmt.default_case {
                    self.resolve_statement(default_case)?;
                }
            }
            Statement::Recover(recover_stmt) => {
                for arm in &mut recover_stmt.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.resolve_expression(guard)?;
                    }
                    self.resolve_statement(&mut arm.body)?;
                }
            }
            Statement::Export(_) => {}
            Statement::Break(_) | Statement::Continue(_) | Statement::Empty => {}
        }
        Ok(())
    }

    fn resolve_expression(&mut self, expression: &mut Expression) -> Result<Option<Type>> {
        match expression {
            Expression::Literal(lit) => Ok(Some(self.literal_type(lit))),
            Expression::Variable(name) => Ok(self.lookup_variable_type(name)),
            Expression::Call(call) => self.resolve_call(call),
            Expression::MethodCall(method_call) => {
                self.resolve_expression(&mut method_call.object)?;
                for arg in &mut method_call.arguments {
                    self.resolve_expression(arg)?;
                }
                Ok(None)
            }
            Expression::Binary(binary) => {
                self.resolve_expression(&mut binary.left)?;
                self.resolve_expression(&mut binary.right)?;
                match binary.operator {
                    crate::ast::expr::BinaryOp::Equal
                    | crate::ast::expr::BinaryOp::NotEqual
                    | crate::ast::expr::BinaryOp::Less
                    | crate::ast::expr::BinaryOp::LessEqual
                    | crate::ast::expr::BinaryOp::Greater
                    | crate::ast::expr::BinaryOp::GreaterEqual
                    | crate::ast::expr::BinaryOp::LogicalAnd
                    | crate::ast::expr::BinaryOp::LogicalOr => {
                        Ok(Some(Type::Basic(crate::ast::types::BasicType::Bool)))
                    }
                    _ => Ok(None),
                }
            }
            Expression::Unary(unary) => {
                self.resolve_expression(&mut unary.operand)?;
                Ok(None)
            }
            Expression::FieldAccess(field_access) => {
                self.resolve_expression(&mut field_access.object)?;
                Ok(None)
            }
            Expression::IndexAccess(index_access) => {
                self.resolve_expression(&mut index_access.object)?;
                self.resolve_expression(&mut index_access.index)?;
                Ok(None)
            }
            Expression::SliceRange(slice_range) => {
                self.resolve_expression(&mut slice_range.object)?;
                if let Some(start) = &mut slice_range.start {
                    self.resolve_expression(start)?;
                }
                if let Some(end) = &mut slice_range.end {
                    self.resolve_expression(end)?;
                }
                Ok(None)
            }
            Expression::Array(array_expr) => {
                for element in &mut array_expr.elements {
                    self.resolve_expression(element)?;
                }
                Ok(None)
            }
            Expression::Slice(slice_expr) => {
                for element in &mut slice_expr.elements {
                    self.resolve_expression(element)?;
                }
                Ok(None)
            }
            Expression::Map(map_expr) => {
                for entry in &mut map_expr.entries {
                    self.resolve_expression(&mut entry.key)?;
                    self.resolve_expression(&mut entry.value)?;
                }
                Ok(None)
            }
            Expression::Struct(struct_expr) => {
                for field in &mut struct_expr.fields {
                    self.resolve_expression(&mut field.value)?;
                }
                Ok(Some(struct_expr.struct_type.clone()))
            }
            Expression::Tuple(tuple_expr) => {
                for element in &mut tuple_expr.elements {
                    self.resolve_expression(element)?;
                }
                Ok(None)
            }
            Expression::If(if_expr) => {
                self.resolve_expression(&mut if_expr.condition)?;
                self.resolve_expression(&mut if_expr.then_branch)?;
                if let Some(else_branch) = &mut if_expr.else_branch {
                    self.resolve_expression(else_branch)?;
                }
                Ok(None)
            }
            Expression::Match(match_expr) => {
                self.resolve_expression(&mut match_expr.expr)?;
                for arm in &mut match_expr.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.resolve_expression(guard)?;
                    }
                    self.resolve_expression(&mut arm.body)?;
                }
                Ok(None)
            }
            Expression::Lambda(lambda) => {
                let saved_env = self.variable_env.clone();
                self.enter_scope();
                for param in &lambda.parameters {
                    self.bind_variable(&param.name, param.param_type.clone());
                }
                let body_type = self.resolve_expression(&mut lambda.body)?;
                self.variable_env = saved_env;
                Ok(body_type.or_else(|| lambda.return_type.clone()))
            }
            Expression::FunctionLiteral(func_lit) => {
                if let Some(lambda) = &mut func_lit.lambda {
                    let saved_env = self.variable_env.clone();
                    self.enter_scope();
                    for param in &lambda.parameters {
                        self.bind_variable(&param.name, param.param_type.clone());
                    }
                    self.resolve_expression(&mut lambda.body)?;
                    self.variable_env = saved_env;
                }
                Ok(func_lit.function_type.clone())
            }
            Expression::Block(block_expr) => {
                self.enter_scope();
                for stmt in &mut block_expr.statements {
                    self.resolve_statement(stmt)?;
                }
                let result = if let Some(final_expr) = &mut block_expr.final_expr {
                    self.resolve_expression(final_expr)?
                } else {
                    None
                };
                self.exit_scope();
                Ok(result)
            }
            Expression::Cast(cast_expr) => {
                self.resolve_expression(&mut cast_expr.expr)?;
                Ok(Some(cast_expr.target_type.clone()))
            }
            Expression::Assert(assert_expr) => {
                self.resolve_expression(&mut assert_expr.expr)?;
                Ok(None)
            }
            Expression::New(new_expr) => {
                if let Some(initializer) = &mut new_expr.initializer {
                    self.resolve_expression(initializer)?;
                }
                Ok(None)
            }
            Expression::Rc(rc_expr) => {
                if let Some(initializer) = &mut rc_expr.initializer {
                    self.resolve_expression(initializer)?;
                }
                Ok(None)
            }
            Expression::Path(_) | Expression::Panic(_) | Expression::Go(_) => Ok(None),
            _ => Ok(None),
        }
    }

    fn resolve_call(&mut self, call: &mut CallExpr) -> Result<Option<Type>> {
        self.resolve_expression(&mut call.callee)?;
        let mut argument_types = Vec::new();
        let mut has_unknown = false;
        for arg in &mut call.arguments {
            match self.resolve_expression(arg)? {
                Some(arg_type) => argument_types.push(arg_type),
                None => {
                    has_unknown = true;
                }
            }
        }

        if has_unknown {
            return Ok(None);
        }

        if let Expression::Variable(name) = &*call.callee {
            let call_type_arg_count = call.type_args.len();
            if let Some(Type::Function(_)) = self.lookup_variable_type(name) {
                // Calling through a function-valued variable; overload resolution not needed
                return Ok(None);
            }
            let candidates = self.symbol_table.lookup_function_overloads(name);
            let is_overloaded_name = candidates.len() > 1;
            if candidates.is_empty() {
                return Err(CompilerError::semantic(
                    call.location.line,
                    call.location.column,
                    format!("Undefined function '{}'", name),
                ));
            }

            let mut matches = Vec::new();
            for candidate in candidates.iter() {
                let candidate = *candidate;
                let candidate_generic_count = candidate.declaration.generics.len();
                if call_type_arg_count > 0 && candidate_generic_count == 0 {
                    // Function is not generic but call provides type arguments
                    continue;
                }
                if candidate_generic_count > 0 && call_type_arg_count != candidate_generic_count {
                    // Generic function requires matching number of type arguments
                    continue;
                }

                let params = &candidate.signature.parameter_types;
                let is_variadic = candidate.signature.is_variadic;
                if (!is_variadic && params.len() != argument_types.len())
                    || (is_variadic && argument_types.len() < params.len())
                {
                    continue;
                }
                let mut compatible = true;
                for (expected, actual) in params.iter().zip(argument_types.iter()) {
                    let (expected_cmp, treat_as_generic) =
                        self.normalize_expected_type(expected, &candidate.declaration);
                    if treat_as_generic {
                        continue;
                    }
                    // BasicType::Any acts as a wildcard to support generic-like builtins.
                    if matches!(
                        expected_cmp,
                        Type::Basic(crate::ast::types::BasicType::Any)
                    ) {
                        continue;
                    }
                    if mangling::type_tag(&expected_cmp) != mangling::type_tag(actual) {
                        if self.bool_int_compatible(&expected_cmp, actual) {
                            continue;
                        }
                        compatible = false;
                        break;
                    }
                }
                if compatible {
                    matches.push(candidate);
                }
            }

            if matches.is_empty() {
                return Err(CompilerError::semantic(
                    call.location.line,
                    call.location.column,
                    format!(
                        "No overload of '{}' matches the provided argument types",
                        name
                    ),
                ));
            }

            if matches.len() > 1 {
                return Err(CompilerError::semantic(
                    call.location.line,
                    call.location.column,
                    format!("Ambiguous call to '{}'", name),
                ));
            }

            let resolved = matches[0];
            let mangled_name = if resolved.declaration.body.is_some() && is_overloaded_name {
                mangling::mangle_function_name(&resolved.signature)
            } else {
                resolved.name.clone()
            };
            call.resolved = Some(ResolvedFunction {
                signature: resolved.signature.clone(),
                mangled_name,
            });

            let return_type = resolved
                .declaration
                .return_type
                .clone()
                .unwrap_or(Type::Basic(crate::ast::types::BasicType::Void));
            Ok(Some(return_type))
        } else {
            Ok(None)
        }
    }

    fn normalize_expected_type(&self, expected: &Type, decl: &FunctionDecl) -> (Type, bool) {
        if let Type::Generic(name) = expected {
            let is_decl_generic = decl.generics.iter().any(|g| g.name == *name);
            if is_decl_generic {
                return (expected.clone(), true);
            }
        }

        let mut normalized = expected.clone();
        self.rewrite_named_generics(&mut normalized);
        (normalized, false)
    }

    fn lookup_named_type(&self, name: &str) -> Option<Type> {
        if self.symbol_table.lookup_struct(name).is_some() {
            return Some(Type::Struct(crate::ast::types::StructType {
                name: name.to_string(),
                type_args: vec![],
                location: crate::error::Location::new(0, 0, 0),
            }));
        }
        if self.symbol_table.lookup_enum(name).is_some() {
            return Some(Type::Enum(crate::ast::types::EnumType {
                name: name.to_string(),
                type_args: vec![],
                location: crate::error::Location::new(0, 0, 0),
            }));
        }
        if self.symbol_table.lookup_type(name).is_some() {
            return Some(Type::Alias(crate::ast::types::AliasType {
                name: name.to_string(),
                type_args: vec![],
                location: crate::error::Location::new(0, 0, 0),
            }));
        }
        None
    }

    fn bool_int_compatible(&self, expected: &Type, actual: &Type) -> bool {
        matches!(expected, Type::Basic(crate::ast::types::BasicType::Bool))
            && matches!(actual, Type::Basic(crate::ast::types::BasicType::Int))
    }

    fn rewrite_named_generics(&self, ty: &mut Type) {
        match ty {
            Type::Generic(name) => {
                if let Some(resolved) = self.lookup_named_type(name) {
                    *ty = resolved;
                }
            }
            Type::Array(array_type) => {
                self.rewrite_named_generics(array_type.element_type.as_mut());
            }
            Type::Map(map_type) => {
                self.rewrite_named_generics(map_type.key_type.as_mut());
                self.rewrite_named_generics(map_type.value_type.as_mut());
            }
            Type::Tuple(tuple_type) => {
                for element in &mut tuple_type.element_types {
                    self.rewrite_named_generics(element);
                }
            }
            Type::Struct(struct_type) => {
                for arg in &mut struct_type.type_args {
                    self.rewrite_named_generics(arg);
                }
            }
            Type::Enum(enum_type) => {
                for arg in &mut enum_type.type_args {
                    self.rewrite_named_generics(arg);
                }
            }
            Type::Function(function_type) => {
                for param in &mut function_type.parameter_types {
                    self.rewrite_named_generics(param);
                }
                if let Some(ret) = &mut function_type.return_type {
                    self.rewrite_named_generics(ret.as_mut());
                }
            }
            Type::Reference(reference_type) => {
                self.rewrite_named_generics(reference_type.referenced_type.as_mut());
            }
            Type::Vec(vec_type) => {
                self.rewrite_named_generics(vec_type.element_type.as_mut());
            }
            Type::Channel(channel_type) => {
                self.rewrite_named_generics(channel_type.element_type.as_mut());
            }
            Type::Optional(optional_type) => {
                self.rewrite_named_generics(optional_type.inner_type.as_mut());
            }
            Type::Error(error_type) => {
                self.rewrite_named_generics(error_type.inner_type.as_mut());
            }
            Type::Result(result_type) => {
                self.rewrite_named_generics(result_type.ok_type.as_mut());
                self.rewrite_named_generics(result_type.err_type.as_mut());
            }
            Type::Option(option_type) => {
                self.rewrite_named_generics(option_type.some_type.as_mut());
            }
            Type::Union(union_type) => {
                for member in &mut union_type.member_types {
                    self.rewrite_named_generics(member);
                }
            }
            Type::Alias(alias_type) => {
                for arg in &mut alias_type.type_args {
                    self.rewrite_named_generics(arg);
                }
            }
            Type::Rc(rc_type) => {
                self.rewrite_named_generics(rc_type.inner_type.as_mut());
            }
            Type::Weak(weak_type) => {
                self.rewrite_named_generics(weak_type.inner_type.as_mut());
            }
            _ => {}
        }
    }

    fn literal_type(&self, literal: &Literal) -> Type {
        match literal {
            Literal::Integer(_) => Type::Basic(crate::ast::types::BasicType::Int),
            Literal::Float(_) => Type::Basic(crate::ast::types::BasicType::Float64),
            Literal::String(_) => Type::Basic(crate::ast::types::BasicType::String),
            Literal::Char(_) => Type::Basic(crate::ast::types::BasicType::Char),
            Literal::Boolean(_) => Type::Basic(crate::ast::types::BasicType::Bool),
            Literal::Null => Type::Basic(crate::ast::types::BasicType::Any),
        }
    }

    fn bind_variable(&mut self, name: &str, ty: Type) {
        if let Some(scope) = self.variable_env.last_mut() {
            scope.insert(name.to_string(), ty);
        }
    }

    fn lookup_variable_type(&self, name: &str) -> Option<Type> {
        for scope in self.variable_env.iter().rev() {
            if let Some(ty) = scope.get(name) {
                let mut normalized = ty.clone();
                self.rewrite_named_generics(&mut normalized);
                return Some(normalized);
            }
        }
        if let Some(var_symbol) = self.symbol_table.lookup_variable(name) {
            if let Some(var_type) = &var_symbol.declaration.var_type {
                let mut normalized = var_type.clone();
                self.rewrite_named_generics(&mut normalized);
                return Some(normalized);
            }
        }
        if let Some(const_symbol) = self.symbol_table.lookup_constant(name) {
            if let Some(const_type) = &const_symbol.declaration.const_type {
                let mut normalized = const_type.clone();
                self.rewrite_named_generics(&mut normalized);
                return Some(normalized);
            }
        }
        None
    }

    fn enter_scope(&mut self) {
        self.variable_env.push(HashMap::new());
    }

    fn exit_scope(&mut self) {
        self.variable_env.pop();
    }
}
